#
#  Copyright 2021 The FATE Authors. All Rights Reserved.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
#

import inspect
import logging
import os
import traceback
import typing
from logging.handlers import TimedRotatingFileHandler
from threading import RLock


class LoggerFactory(object):
    TYPE = "FILE"
    LOG_FORMAT = "[%(levelname)s] [%(asctime)s] [source_job] [%(process)s:%(thread)s] - [%(module)s.%(funcName)s] [line:%(lineno)d]: %(message)s"
    LEVEL = logging.getLevelName(os.getenv("OI_LOG_LEVEL", "INFO"))
    logger_dict = {}
    global_handler_dict = {}

    LOG_DIR = None
    PARENT_LOG_DIR = None
    log_share = True

    append_to_parent_log = None

    lock = RLock()
    # CRITICAL = 50
    # FATAL = CRITICAL
    # ERROR = 40
    # WARNING = 30
    # WARN = WARNING
    # INFO = 20
    # DEBUG = 10
    # NOTSET = 0
    levels = (10, 20, 30, 40)

    @classmethod
    def set_directory(cls, directory=None, parent_log_dir=None, append_to_parent_log=None, force=False):
        if not directory:
            raise RuntimeError("directory can not empty")
        if parent_log_dir:
            cls.PARENT_LOG_DIR = parent_log_dir
        if append_to_parent_log:
            cls.append_to_parent_log = append_to_parent_log
        with cls.lock:
            if not cls.LOG_DIR or force:
                cls.LOG_DIR = directory
            if cls.log_share:
                oldmask = os.umask(000)
                os.makedirs(cls.LOG_DIR, exist_ok=True)
                os.umask(oldmask)
            else:
                os.makedirs(cls.LOG_DIR, exist_ok=True)
            for loggerName, ghandler in cls.global_handler_dict.items():
                for className, (logger, handler) in cls.logger_dict.items():
                    logger.removeHandler(ghandler)
                ghandler.close()
            cls.global_handler_dict = {}
            for className, (logger, handler) in cls.logger_dict.items():
                logger.removeHandler(handler)
                _handler = None
                if handler:
                    handler.close()
                if className != "default":
                    _handler = cls.get_handler(className)
                    logger.addHandler(_handler)
                cls.assemble_global_handler(logger)
                cls.logger_dict[className] = logger, _handler

    @classmethod
    def set_default_log_dir(cls, log_dir: str):
        cls.LOG_DIR = log_dir

    @classmethod
    def get_default_log_dir(cls):
        return cls.LOG_DIR or os.getenv("LOG_DIR", default=None)

    @classmethod
    def get_logger(cls, name=None, append_level_log=True):
        with cls.lock:
            if name in cls.logger_dict.keys():
                logger, handler = cls.logger_dict[name]
                if not logger:
                    logger, handler = cls.init_logger(name, append_level_log)
            else:
                logger, handler = cls.init_logger(name, append_level_log)
            return logger

    @classmethod
    def get_global_handler(cls, logger_name, level=None, log_dir=None):
        if not cls.get_default_log_dir():
            return logging.StreamHandler()
        if log_dir:
            logger_name_key = logger_name + "_" + log_dir
        else:
            logger_name_key = logger_name + "_" + cls.get_default_log_dir()
        if logger_name_key not in cls.global_handler_dict:
            with cls.lock:
                if logger_name_key not in cls.global_handler_dict:
                    handler = cls.get_handler(logger_name, level, log_dir)
                    cls.global_handler_dict[logger_name_key] = handler
        return cls.global_handler_dict[logger_name_key]

    @classmethod
    def get_handler(cls, name, level=None, log_dir=None, log_type=None, job_id=None, suffix=""):
        log_dir = log_dir if log_dir else cls.get_default_log_dir()
        if not log_dir:
            return logging.StreamHandler()
        if not name and not log_type:
            return logging.StreamHandler()
        elif not log_type:
            log_file_prefix = os.path.join(log_dir, name)
        elif not name:
            log_file_prefix = os.path.join(log_dir, log_type)
        else:
            log_file_prefix = os.path.join(log_dir, f"{name}_{log_type}")
        if suffix:
            log_file = f"{log_file_prefix}_{suffix}.log"
        else:
            log_file = f"{log_file_prefix}.log"
        job_id = job_id or os.getenv("JOB_ID")
        if job_id:
            source_job = job_id
            if os.getenv("ROLE"):
                source_job += f":{os.getenv('ROLE').lower()}"
                if os.getenv("NODE_ID"):
                    source_job += f":{os.getenv('NODE_ID')}"
            formatter = logging.Formatter(cls.LOG_FORMAT.replace("source_job", source_job))
        else:
            formatter = logging.Formatter(cls.LOG_FORMAT.replace("source_job", "unknow"))
        os.makedirs(os.path.dirname(log_file), exist_ok=True)
        if cls.log_share:
            handler = ROpenHandler(log_file,
                                   when='D',
                                   interval=1,
                                   backupCount=14,
                                   delay=True)
        else:
            handler = TimedRotatingFileHandler(log_file,
                                               when='D',
                                               interval=1,
                                               backupCount=14,
                                               delay=True)

        if level:
            handler.level = level

        handler.setFormatter(formatter)
        return handler

    @classmethod
    def init_logger(cls, logger_name, append_level_log=True):
        with cls.lock:
            logger = cls.new_logger(logger_name)
            handler = None
            if logger_name:
                handler = cls.get_handler(logger_name)
                logger.addHandler(handler)
                cls.logger_dict[logger_name] = logger, handler

            else:
                cls.logger_dict["default"] = logger, handler

            if append_level_log:
                cls.assemble_global_handler(logger)
            return logger, handler

    @classmethod
    def assemble_global_handler(cls, logger):
        # add level loggers to the peer directory
        if cls.get_default_log_dir():
            for level in cls.levels:
                if level >= cls.LEVEL:
                    level_logger_name = logging._levelToName[level].lower()
                    logger.addHandler(cls.get_global_handler(level_logger_name, level))
        # add level loggers to the parent directory
        if cls.append_to_parent_log and cls.PARENT_LOG_DIR:
            for level in cls.levels:
                if level >= cls.LEVEL:
                    level_logger_name = logging._levelToName[level].lower()
                    logger.addHandler(
                        cls.get_global_handler(level_logger_name, level, cls.PARENT_LOG_DIR))

    @classmethod
    def new_logger(cls, logger_name):
        logger = logging.getLogger(logger_name)
        logger.propagate = False
        logger.setLevel(cls.LEVEL)
        return logger


def setDirectory(directory=None):
    LoggerFactory.set_directory(directory)


def setLevel(level):
    LoggerFactory.LEVEL = level


def getLogger(name=None, append_level_log=True):
    if name is None:
        frame = inspect.stack()[1]
        module = inspect.getmodule(frame[0])
        name = "stats"
    return LoggerFactory.get_logger(name, append_level_log)


def exception_to_trace_string(ex):
    return "".join(traceback.TracebackException.from_exception(ex).format())


class ROpenHandler(TimedRotatingFileHandler):
    def _open(self):
        prevumask = os.umask(000)
        rtv = TimedRotatingFileHandler._open(self)
        os.umask(prevumask)
        return rtv


def ready_log(msg, job=None, task=None, role=None, party_id=None, detail=None):
    prefix, suffix = base_msg(job, task, role, party_id, detail=detail)
    return f"{prefix}{msg} ready{suffix}"


def start_log(msg, job=None, task=None, role=None, party_id=None, detail=None):
    prefix, suffix = base_msg(job, task, role, party_id, detail=detail)
    return f"{prefix}start to {msg}{suffix}"


def successful_log(msg, job=None, task=None, role=None, party_id=None, code=None, detail=None):
    prefix, suffix = base_msg(job, task, role, party_id, code, detail)
    return f"{prefix}{msg} successfully{suffix}"


def warning_log(msg, job=None, task=None, role=None, party_id=None, code=None, detail=None):
    prefix, suffix = base_msg(job, task, role, party_id, code, detail)
    return f"{prefix}{msg} is not effective{suffix}"


def failed_log(msg, job=None, task=None, role=None, party_id=None, code=None, detail=None):
    prefix, suffix = base_msg(job, task, role, party_id, code, detail)
    return f"{prefix}failed to {msg}{suffix}"


def base_msg(job=None, task=None, role: str = None, party_id: typing.Union[str, int] = None, code=None, detail=None):
    if code:
        detail_msg = f" code: {code}"
    else:
        detail_msg = ""
    if detail:
        detail_msg += f" detail: \n{detail}"
    if task is not None:
        return f"task {task.f_task_id} {task.f_task_version} ", f" on {task.f_role} {task.f_party_id}{detail_msg}"
    elif job is not None:
        return "", f" on {job.f_role} {job.f_party_id}{detail_msg}"
    elif role and party_id:
        return "", f" on {role} {party_id}{detail_msg}"
    else:
        return "", f"{detail_msg}"
